#include <bits/stdc++.h>
using namespace std;

/*
3234. 统计 1 显著的字符串的数量
已解答
中等
相关标签
premium lock icon
相关企业
提示
给你一个二进制字符串 s。

请你统计并返回其中 1 显著 的 子字符串 的数量。

如果字符串中 1 的数量 大于或等于 0 的数量的 平方，则认为该字符串是一个 1 显著 的字符串 。

 

示例 1：

输入：s = "00011"

输出：5

解释：

1 显著的子字符串如下表所示。

i	j	s[i..j]	0 的数量	1 的数量
3	3	1	0	1
4	4	1	0	1
2	3	01	1	1
3	4	11	0	2
2	4	011	1	2
示例 2：

输入：s = "101101"

输出：16

解释：

1 不显著的子字符串如下表所示。

总共有 21 个子字符串，其中 5 个是 1 不显著字符串，因此有 16 个 1 显著子字符串。

i	j	s[i..j]	0 的数量	1 的数量
1	1	0	1	0
4	4	0	1	0
1	4	0110	2	2
0	4	10110	2	3
1	5	01101	2	3
 

提示：

1 <= s.length <= 4 * 104
s 仅包含字符 '0' 和 '1'。
*/

// 法一
class Solution {
public:
    int numberOfSubstrings(string s) {
        int n = s.length();
        vector<int> pos(n + 1);
        pos[n] = n; // 边界处理：超出字符串范围的"0位置"设为n
        // 预处理pos数组：pos[i] = i及之后第一个'0'的位置（无0则为n）
        for (int i = n - 1; i >= 0; i--) {
            pos[i] = (s[i] == '0') ? i : pos[i + 1];
        }

        int ans = 0;
        // 枚举所有子串起点l
        for (int l = 0; l < n; l++) {
            int r0 = pos[l]; // l后第一个'0'的位置
            // 情况1：l到末尾全是'1'（0的数量为0，必满足条件）
            if (r0 == n) {
                ans += n - l; // 所有以l为起点的子串都有效，共n-l个
                continue;
            }
            // 情况2：l到第一个'0'前的纯'1'子串（0数=0，有效）
            ans += r0 - l;

            int zero = 0; // 当前子串中'0'的个数
            int current_r0 = r0; // 当前遍历到的'0'的位置
            // 按'0'的个数分组统计（zero递增，直到zero²超过最大可能1数）
            while (current_r0 < n) {
                zero++; // 新增一个'0'
                // 剪枝：zero²超过l到末尾的总长度（最大可能1数≤总长度），后续无效
                if (zero * zero > n - l) break;

                int r1 = pos[current_r0 + 1]; // 下一个'0'的位置（或n）
                int len = current_r0 - l + 1; // 当前子串[l..current_r0]的长度
                int cnt1 = len - zero; // 当前子串中'1'的数量
                // 满足条件需要的额外'1'数（不够则补0）
                int need = max(zero * zero - cnt1, 0);
                // 统计当前zero数下有效子串数：[current_r0+need .. r1-1]区间的子串都有效
                ans += max(r1 - current_r0 - need, 0);
                // 移动到下一个'0'的位置，继续分组
                current_r0 = r1;
            }
        }
        return ans;
    }
};

// 法二
class Solution {
public:
    int numberOfSubstrings(string s) {
        int res = 0;          // 结果：1显著子串总数
        int n = s.size();
        vector<int> pos{-1};  // 存储所有'0'的位置（开头加-1作为左边界）
        int total1 = 0;       // 截至当前终点r的总1数量

        // 枚举所有子串终点r
        for(int r = 0; r < n; r++) {
            if(s[r] == '0')  pos.push_back(r); // 记录当前0的位置
            else {
                total1++; // 累计1的数量
                // 统计：以r为终点、0数量为0的有效子串（纯1子串）
                res += r - pos.back();
            }
            int m = pos.size(); // m = 0的个数 + 1（左边界）
            // 内层循环：按0的个数cnt0分组，统计有效子串（cnt0从大到小）
            // 条件：cnt0² <= total1（否则不可能满足1显著条件）
            for(int i = m - 1; i > 0 && (m - i) * (m - i) <= total1; i--) {
                int p = pos[i-1]; // 上一个0的位置（左边界）
                int q = pos[i];   // 当前0的位置
                int cnt0 = m - i; // 当前分组的0数量（pos[i..m-1]共cnt0个0）
                // 当前区间[q, r]中1的数量
                int cnt1 = (r - q + 1) - cnt0;
                // 满足条件需补充的1数，不够则为0
                int need = max(0, cnt0 * cnt0 - cnt1);
                // 统计当前分组的有效子串数：[p+1, q-need]区间的起点都有效
                res += max(0, q - need - p);
            }
        }

        return res;
    }
};